from pymongo import MongoClient, ASCENDING, DESCENDING
from pymongo.errors import ConnectionFailure
from config import config
import logging

logger = logging.getLogger(__name__)


class Database:
    """数据库管理类"""
    
    def __init__(self):
        self.mongo_client = None
        self.db = None
    
    def connect_mongodb(self):
        """连接MongoDB"""
        try:
            self.mongo_client = MongoClient(
                config.MONGODB_URI,
                serverSelectionTimeoutMS=5000
            )
            # 测试连接
            self.mongo_client.admin.command('ping')
            self.db = self.mongo_client[config.MONGODB_DB]
            logger.info(f"成功连接到MongoDB: {config.MONGODB_DB}")
            
            # 创建索引
            self._create_indexes()
            return True
        except ConnectionFailure as e:
            logger.error(f"MongoDB连接失败: {e}")
            return False
        except Exception as e:
            logger.error(f"MongoDB连接异常: {e}")
            return False
    
    def _create_indexes(self):
        """创建数据库索引"""
        try:
            # 设备表索引
            self.db.devices.create_index([("deviceId", ASCENDING)], unique=True)
            self.db.devices.create_index([("status", ASCENDING)])
            self.db.devices.create_index([("lastHeartbeatTime", DESCENDING)])
            
            # 心跳日志索引
            self.db.heartbeat_logs.create_index([("deviceId", ASCENDING)])
            self.db.heartbeat_logs.create_index([("timestamp", DESCENDING)])
            self.db.heartbeat_logs.create_index([
                ("deviceId", ASCENDING),
                ("timestamp", DESCENDING)
            ])
            
            # 设备日志索引
            self.db.device_logs.create_index([("deviceId", ASCENDING)])
            self.db.device_logs.create_index([("timestamp", DESCENDING)])
            self.db.device_logs.create_index([("eventType", ASCENDING)])
            
            logger.info("数据库索引创建完成")
        except Exception as e:
            logger.error(f"创建索引失败: {e}")
    
    def close(self):
        """关闭数据库连接"""
        if self.mongo_client:
            self.mongo_client.close()
            logger.info("MongoDB连接已关闭")


# 全局数据库实例
db = Database()

